# sa.py - state recording and analysis

import sys
import stackless
import pickle
import inspect
from objdump import objdump
import dis
import inspect

import sys
from cStringIO import StringIO
import sqlite3

import copy

import cPickle

from types import *
import multiprocessing
import time


old_stdout = sys.stdout
dump_output = StringIO()
sys.stdout = dump_output

try:
	from fib import *
except ImportError:
	pass

sys.stdout = old_stdout


import subprocess


# Configuration

# If fold_lines is True we record a single state for each line executed,
# not each bytecode.


fold_lines = False


CSS = 0
PSS = 1
record_mode = CSS

# One thing is very weird: While using parallel in query_expr
# I will always get the same performance as single Process at first time
# But better performance after that
# This will not happen to single Process, no matter how many times I run
parallel_mode = False


# record_modules lists the modules for which we want to record execution.
record_modules = ['__main__', 'external_module']

# Notes
# - Only supports single-threaded applications.
#   - To handle multiple threads we need to modifiy the thread module

def SA_ROOT(filename):
    """Root function for tasket container"""
    # Setup fake global name space
    sa_root_globals = {}
    sa_root_globals['__builtins__'] = globals()['__builtins__']
    sa_root_globals['__name__'] = '__main__'
    sa_root_globals['__doc__'] = None
    sa_root_globals['__package__'] = None

#   Comment these 'print' statements, since they will also be counted in __main__
#    print '[state] root begin'

    # Ensure that program appears to be running at the top level namespace.
    execfile(filename, sa_root_globals, sa_root_globals)

#    print '[state] root end'




class SA_Control(object):
    """State Analysis Control Container"""

    def __init__(self):
        self.state_count = 0
        self.state_skips = 0


    def start(self, modulename, filename, sfilename, record=False):
        self.filename = filename
        self.sfilename = sfilename
        self.states = []
        self.state_count = 0
        self.state_skips = 0


#        print "execute it here already"
#
#        exec "from count import *"

        self.task = stackless.tasklet( SA_ROOT )(self.filename)

        if record:
            # redirect console output to a string first
            self.old_stdout = sys.stdout
            self.console_output = StringIO()
            sys.stdout = self.console_output


    def init_database(self):

        conn = sqlite3.connect(self.sfilename)
        c = conn.cursor()

        try:
            c.execute('''create table console (id INTEGER, sindex INTEGER, console_msg text)''')
        except sqlite3.OperationalError:
            c.execute('''delete from console''')

        try:
            c.execute('''create table source_file (sindex INTEGER, code text)''')
        except sqlite3.OperationalError:
            c.execute('''delete from console''')


        try:
            if record_mode == CSS:
                c.execute('''create table states (id INTEGER PRIMARY KEY ASC, co_name text, module_name text, f_globals blob, f_locals blob, source_line_no INTEGER)''')
            elif record_mode == PSS:
                c.execute('''create table states (id INTEGER PRIMARY KEY ASC, state blob)''')
        except sqlite3.OperationalError:
            c.execute('''drop table states''')

            if record_mode == CSS:
                c.execute('''create table states (id INTEGER PRIMARY KEY ASC, co_name text, module_name text, f_globals blob, f_locals blob, source_line_no INTEGER)''')
            elif record_mode == PSS:
                c.execute('''create table states (id INTEGER PRIMARY KEY ASC, state blob)''')






        conn.commit()
        c.close()



    def load(self, sfilename, index=0):

        self.sfilename = sfilename
        self.conn = sqlite3.connect(self.sfilename)
        self.conn.row_factory = sqlite3.Row
        self.c = self.conn.cursor()

        if record_mode == PSS:
            self.c.execute("select state from states where id=%d" %(index))
            task = str(self.c.fetchone()[0])
            self.task = cPickle.loads(task)

            self.task.insert()
            self.state_count = index


    def save_console_table(self):
        sys.stdout = self.old_stdout

        conn = sqlite3.connect(self.sfilename)
        conn.isolation_level = "DEFERRED"
        c = conn.cursor()

        output = {}
        lines = self.console_output.getvalue().split("\n")

        index = 0
        for line in lines:
            if line.startswith('**********'):
                if index not in output:
                    output[index] = ""
                index = int(line.split()[1])

            else:
                if index not in output:
                    output[index] = line
                else:
                    output[index] += line

        keys = output.keys()
        keys.sort()

        index = 0
        for key in keys:
            if output[key] != "":
                print output[key]
                c.execute('''insert into console values ('%d', '%d', "%s")''' %(index, key, output[key]))
                index += 1


        conn.commit()
        c.close()



    def get_state_size(self):

        if "c" not in self.__dict__:
            self.load("record.pys", 0)

        self.c.execute('select max(id) from states')

        state_size = self.c.fetchone()[0]

        return state_size



    def play(self):
        self.c.execute('select console_msg from console where sindex >= %d' %(int(self.state_count)))

        result = ""

        for row in self.c:
            result += "%s\n" %(row['console_msg'])

        print result


    def end(self):

        print '[state] state count = %d' % (self.state_count)
        print '[state] state skips = %d' % (self.state_skips)



    def step(self, count=1):
        self.task.insert()
        self.task = stackless.run(count)
        return self.task


    def eval_expr(self, expression, globals, locals):

        try:
            if eval(expression, globals, locals):
                return True
        except:
            pass

        return False


    def parallel_query_expr(self, start_index, size, expression, PROCESSES):

        result_index = {}

        conn = sqlite3.connect(self.sfilename)
        p = conn.cursor()

        id = int(multiprocessing.current_process().name[-1:])

        for i in range(id%PROCESSES):
            conn = sqlite3.connect(self.sfilename)
            p = conn.cursor()


        p.execute('select id, f_globals, f_locals from states limit %d, %d' %(start_index, size))


        for row in p:
            state_index = row[0]

            f_globals = cPickle.loads(str(row[1]))

            f_locals = cPickle.loads(str(row[2]))

            if self.eval_expr(expression, f_globals, f_locals):
                result_index[state_index] = True


        return result_index


    def calculate(self, func, args):
        result = func(*args)
        return result


    def query_expr(self, expression):

        result = {}

        if record_mode == CSS:

            PROCESSES = multiprocessing.cpu_count()

            if PROCESSES == 1 or not parallel_mode:

                self.c.execute("select id, co_name, f_globals, f_locals from states")


                for row in self.c:

                    state_index = row[0]

                    code_name = str(row[1])

                    f_globals = cPickle.loads(str(row[2]))

                    f_locals = cPickle.loads(str(row[3]))

                    if code_name == '<module>':

                        if self.eval_expr(expression, f_globals, f_locals):
                            result[state_index] = True


                    else:
                        if 'self' in f_locals:
                            if type(f_locals['self']) == InstanceType:
                                if self.eval_expr(expression, f_globals, f_locals):
                                    result[state_index] = True

                        else:
                            if self.eval_expr(expression, f_globals, f_locals):
                                result[state_index] = True


            else:

                self.c.execute('select max(id) from states')

                max_state_index = int(self.c.fetchone()[0])


                start_index_stack = []

                size = max_state_index / PROCESSES
                for i in range(PROCESSES):
                    start_index_stack.append(i * size)

                pool = multiprocessing.Pool(PROCESSES)

                TASKS = [(self.parallel_query_expr, (i, size + 2, expression, PROCESSES)) for i in start_index_stack]

                parallel_results = [pool.apply_async(self.calculate, t) for t in TASKS]


                for r in parallel_results:
                    result.update(r.get())


        elif record_mode == PSS:

            while self.step(1):

                state_index = self.task.frame.f_lineno

                code_name = self.task.frame.f_code.co_name

                f_globals = self.task.frame.f_globals

                f_locals = self.task.frame.f_locals


                if code_name == '<module>':

                    if self.eval_expr(expression, f_globals, f_locals):
                        result.append(state_index)


                else:
                    if 'self' in f_locals:
                        if type(f_locals['self']) == InstanceType:
                            if self.eval_expr(expression, f_globals, f_locals):
                                result[state_index] = True

                    else:
                        if self.eval_expr(expression, f_globals, f_locals):
                            result[state_index] = True


        return result



    def query_module(self, module_name):

        result = {}

        if record_mode == CSS:

            self.c.execute("select id from states where module_name='%s'" %(module_name))

            for row in self.c:
                result[row['id']] = True


        elif record_mode == PSS:

            while self.step(1):

                if self.task.frame.f_globals['__name__'] == module_name:
                    result[self.task.frame.f_lineno] = True

        return result



    def query_class(self, class_name):

        result = {}

        if record_mode == CSS:

            self.c.execute("select id, co_name, f_globals, f_locals from states where module_name!='__name__'")


            for row in self.c:

                state_index = row[0]

                code_name = str(row[1])

                f_globals = cPickle.loads(str(row[2]))

                f_locals = cPickle.loads(str(row[3]))


                if 'self' in f_locals:
                    if type(f_locals['self']) == InstanceType:
                        obj = f_locals['self']
                        attr = getattr(obj, code_name)

                        if attr.im_class.__name__ == class_name:
                            result[state_index] = True



        elif record_mode == PSS:

            while self.step(1):

                code_name = self.task.frame.f_code.co_name
                if code_name == '<module>':
                    continue


                if 'self' in self.task.frame.f_locals:
                    if type(self.task.frame.f_locals['self']) == InstanceType:
                        obj = self.task.frame.f_locals['self']
                        attr = getattr(obj, code_name)

                        if attr.im_class.__name__ == class_name:
                            result[self.task.frame.f_lineno] = True

        return result



    def query_function(self, function_name):

        result = {}

        if record_mode == CSS:

            self.c.execute("select id, co_name, f_globals from states where module_name!='__name__'")

            for row in self.c:

                state_index = row[0]

                code_name = str(row[1])

                f_globals = cPickle.loads(str(row[2]))

                if code_name == function_name and function_name in f_globals:

                    code = f_globals[function_name]

                    if type(code) == FunctionType:

                        result[state_index] = True


        elif record_mode == PSS:

            while self.step(1):

                if self.task.frame.f_code.co_name == function_name and function_name in self.task.frame.f_globals:

                    code = self.task.frame.f_globals[function_name]

                    if type(code) == FunctionType:

                        result[self.task.frame.f_lineno] = True


        return result



    def query_method(self, method_name):

        result = {}

        if record_mode == CSS:

            self.c.execute("select id, co_name, f_globals, f_locals from states where module_name!='__name__'")

            for row in self.c:

                state_index = row[0]

                code_name = str(row[1])

                f_globals = cPickle.loads(str(row[2]))

                if code_name == method_name:

                    f_locals = cPickle.loads(str(row[3]))

                    if 'self' in f_locals:

                        if type(f_locals['self']) == InstanceType:
                            obj = f_locals['self']
                            attr = getattr(obj, code_name)

                            if type(attr) == MethodType:
                                result[state_index] = True


        elif record_mode == PSS:

            while self.step(1):

                if self.task.frame.f_code.co_name == method_name:

                    if 'self' in self.task.frame.f_locals:

                        if type(self.task.frame.f_locals['self']) == InstanceType:
                            obj = self.task.frame.f_locals['self']
                            attr = getattr(obj, method_name)

                            if type(attr) == MethodType:
                                result[self.task.frame.f_lineno] = True


        return result


    def execute_code(self, state_dict, ns, record=False):

        prev_lineno = 0

        index = 0


        while self.step(1):

            # Skip SA_ROOT
            if self.task.frame.f_code.co_name == "SA_ROOT":
                continue

            # Skip modules
            if self.task.frame.f_globals['__name__'] not in record_modules:
                continue


            # Fold lines

            # one problem is: if the variable is local variable, it might not be able to save
            #     since it has already record the first state

            if fold_lines:
                if self.task.frame.f_lineno == prev_lineno:
                    continue
                else:
                    prev_lineno = self.task.frame.f_lineno

            self.state_count += 1

            if record:

                try:
                    state = None
                    if record_mode == PSS:
                        state = cPickle.dumps(self.task, 2)
                    elif record_mode == CSS:

                        new_globals = {}
                        for k, v in self.task.frame.f_globals.items():
                            if type(v) not in (DictType, DictionaryType):
                                new_globals[k] = copy.deepcopy(v)
                            else:
                                new_globals[k] = v


                        state = [copy.deepcopy(self.task.frame.f_code.co_name),
                                 copy.deepcopy(self.task.frame.f_globals['__name__']),
                                 new_globals,
                                 copy.deepcopy(self.task.frame.f_locals)]

                    state_dict[index] = state

                    index += 1

                    # it is weird, I need \n here, otherwise it will stay in the same line
                    print '\n','*' * 10, self.state_count
                except:
                    self.state_skips += 1


        ns.count = self.state_count

        ns.skips = self.state_skips


        state_dict['STOP'] = True

        self.save_console_table()



    def insert_database(self, state_dict):

        index = 0

        conn = sqlite3.connect(self.sfilename)
        conn.isolation_level = "DEFERRED"
        c = conn.cursor()

        while len(state_dict) != 0 or 'STOP' not in state_dict:
            for key in state_dict.keys():

                if key == 'STOP':
                    if len(state_dict) == 1:

                        conn.commit()
                        c.close()

                        return
                    else:
                        continue

                state = state_dict.pop(key)

                if record_mode == CSS:
                    co_name = state[0]
                    module_name = state[1]
                    f_globals = cPickle.dumps(state[2], 2)
                    f_locals = cPickle.dumps(state[3], 2)

                    print ('''insert into states values(?, ?, ?, ?, ?)''', (index, co_name, module_name, sqlite3.Binary(f_globals), sqlite3.Binary(f_locals)))

                    c.execute('''insert into states values(?, ?, ?, ?, ?)''', (index, co_name, module_name, sqlite3.Binary(f_globals), sqlite3.Binary(f_locals)))
                    index += 1


                elif record_mode == PSS:

                    c.execute('''insert into states values(?, ?)''', (index, sqlite3.Binary(state)))
                    index += 1




    def run(self, record=False):

        prev_lineno = 0

        PROCESSES = multiprocessing.cpu_count()

        if PROCESSES == 1 or not parallel_mode:

            while self.step(1):

                # Skip SA_ROOT
                if self.task.frame.f_code.co_name == "SA_ROOT":
                    continue

                # Skip modules
                if self.task.frame.f_globals['__name__'] not in record_modules:
                    continue


                # Fold lines

                # one problem is: if the variable is local variable, it might not be able to save
                #     since it has already record the first state

                if fold_lines:
                    if self.task.frame.f_lineno == prev_lineno:
                        continue
                    else:
                        prev_lineno = self.task.frame.f_lineno

                self.state_count += 1

                if record:

                    state = None
                    if record_mode == PSS:
                        state = cPickle.dumps(self.task, 2)
                    elif record_mode == CSS:

                        new_globals = {}
                        for k, v in self.task.frame.f_globals.items():
                            if type(v) not in (DictType, DictionaryType):
                                new_globals[k] = v#copy.deepcopy(v)
                            else:
                                new_globals[k] = v

                        new_locals = {}
                        for k, v in self.task.frame.f_locals.items():
                            if type(v) in (IntType, LongType, FloatType, BooleanType, ComplexType, StringTypes, UnicodeType, TupleType, ListType, ClassType, FunctionType):
#                                print type(v)
#                                new_locals[k] = copy.deepcopy(v)
                                new_locals[k] = v
                            else:
                                new_locals[k] = v

                                
                        state = [copy.deepcopy(self.task.frame.f_code.co_name),
                                 copy.deepcopy(self.task.frame.f_globals['__name__']),
                                 new_globals,
                                 new_locals,
                                 self.task.frame.f_lineno - 1]

                    self.states.append(state)


                        # it is weird, I need \n here, otherwise it will stay in the same line
                        # I guess something wrong with the stdout
                    print '\n','*' * 10, self.state_count


            self.save_states_table()

            self.save_console_table()

            self.save_source_file()

        else:

            manager = multiprocessing.Manager()

            # this is a shared dictionary between two Processes
            state_dict = manager.dict()

            # this is shared memory between Process-p and main-Process
            ns = manager.Namespace()
            ns.count = 0
            ns.skips = 0


            p = multiprocessing.Process(target=self.execute_code, args=[state_dict, ns, record])

            q = multiprocessing.Process(target=self.insert_database, args=[state_dict])

            p.start()

            q.start()

            p.join()

            q.join()

            # need to set it again, since anything in Process-p(except for database operations) has gone
            sys.stdout = self.old_stdout

            self.state_count = ns.count
            self.state_skips = ns.skips


    def save_source_file(self):

        conn = sqlite3.connect(self.sfilename)
        c = conn.cursor()

        f = open(self.filename, 'r')

        lines = f.readlines()

        index = 0
        for line in lines:
            c.execute('''insert into source_file values(?, ?)''', (index, line))
            index += 1

        conn.commit()
        c.close()



    def save_states_table(self):

        conn = sqlite3.connect(self.sfilename)
        c = conn.cursor()

        index = 0

        if record_mode == CSS:
            for s in self.states:
                co_name = s[0]
                module_name = s[1]
#                f_globals = cPickle.dumps(s[2], 2)
#                f_locals = cPickle.dumps(s[3], 2)
                f_globals = pickle.dumps(s[2], 2)
                f_locals = pickle.dumps(s[3], 2)
#                print 'local: ', s[3]

                c.execute('''insert into states values(?, ?, ?, ?, ?, ?)''', (index, co_name, module_name, sqlite3.Binary(f_globals), sqlite3.Binary(f_locals), s[4]))
                index += 1

        elif record_mode == PSS:
            for s in self.states:
                c.execute('''insert into states values(?, ?)''', (index, sqlite3.Binary(s)))
                index += 1

        conn.commit()
        c.close()




    def info(self):
        self.c.execute("select max(id) from states")
        states_size = int(str(self.c.fetchone()[0])) + 1

        print '[state] file = %s' % (self.sfilename)
        print '[state] state count = %d' % (states_size)


    def examine(self, ex_code):
        #print self.task.frame.f_locals
        #print 'ex_code = %s' % ex_code
        #objdump(self.task.frame)
        print '[state] line = %d' % (self.task.frame.f_lineno)
        #print self.task.frame.f_globals
        try:
            rv = eval(ex_code, self.task.frame.f_globals, self.task.frame.f_locals)
        except:
            rv = None

        return rv


    def examine_all(self, ex_code):
        for state in self.states:
            self.task = pickle.loads(state)
            rv = self.examine(ex_code)
            print rv


    def read_state(self, id):

        self.c.execute("select f_globals, f_locals from states where id='%s'" %(id))

        result = self.c.fetchone()

        f_globals = cPickle.loads(str(result[0]))

        f_locals = cPickle.loads(str(result[1]))

        return f_globals, f_locals


    def get_source_code(self):


        source_output = ""

        self.c.execute("select code from source_file")

        for row in self.c:
            source_output += str(row[0])


        return source_output


        
    def get_console(self, id=-1):

        conn = sqlite3.connect(self.sfilename)
        c = conn.cursor()

        console_output = ""

        if id == -1:

            c.execute("select console_msg from console")

            for row in c:
                console_output += str(row[0]) + "\n"

        else:

            c.execute("select id from console where sindex=%d" %(id))

            for row in c:
                console_output = str(row[0])

        print 'console_output: ', console_output 


        conn.commit()
        c.close()

        return console_output



    def get_source_id(self, id):

        conn = sqlite3.connect(self.sfilename)
        c = conn.cursor()

        source_id = -1

        print "select source_line_no from states where id=%d" %(id)

        c.execute("select source_line_no from states where id=%d" %(id))

        for row in c:
            source_id = int(row[0])

        return source_id




    def get_console_id(self, id):

        conn = sqlite3.connect(self.sfilename)
        c = conn.cursor()

        console_id = -1

        c.execute("select id from console where sindex=%d" %(id))

        for row in c:
            console_id = int(row[0])

        return console_id


    def append_args(self, args):

        if len(sys.argv) < 2:
            sys.argv.append(args)
        else:
            sys.argv[1] = args

        
    def query_parse(self, expr):

        expr_list = [expr]
        if " or " in expr:
            expr_list = expr.split(" or ")


        state_id_list = [-1]

        for expr_or_item in expr_list:

            or_item_list = []

            if " and " in expr_or_item:
                expr_and_list = expr_or_item.split(" and ")

                and_list = [-1]

                for expr_and_item in expr_and_list:

                    item_list = self.query_driver(expr_and_item)

                    if and_list == [-1]:
                        and_list = item_list
                    else:
                        and_list = [i for i in and_list if i in item_list]

                or_item_list = and_list

            else:

                or_item_list = self.query_driver(expr_or_item)


            if state_id_list == [-1]:
                state_id_list = or_item_list

            else:
                state_id_list += [i for i in or_item_list if i not in state_id_list]


        print "state_id_list:", state_id_list

        return state_id_list

            



    def query_driver(self, expr):

        item_list = []

        if expr.startswith("method "):

            item_list = self.query_method(expr["method ":])

        elif expr.startswith("class "):

            item_list = self.query_method(expr["class ":])

        elif expr.startswith("module "):

            item_list = self.query_module(expr["module ":])

        elif expr.startswith("function "):

            item_list = self.query_function(expr["function ":])

        else:

            item_list = self.query_expr(expr)

        return item_list.keys()
